import os
import shutil
import matplotlib.pyplot as plt

from glob import glob
from utils import get_xml_points, get_map_from_points, get_map_num, get_fidtm_map_num, save_density


# 测试数据、方法信息

data_name = 'DroneRGBT'

gt_path = "../datasets/DroneRGBT/Test/GT_"
rgb_path = "../datasets/DroneRGBT/Test/RGB"
tir_path = "../datasets/DroneRGBT/Test/Infrared"
gt_list = glob(os.path.join(gt_path, '*.xml'))

method_results = {
    'CAGNet': f"../other_methods/MIANet/map_{data_name}",
    'MC3Net': f"../other_methods/MC3Net/map_{data_name}",
    'GETANet': f"../other_methods/GETANet/map_{data_name}",
    'DEFNet': f"../other_methods/DEFNet/map_{data_name}",
    'MSDTrans': f"../other_methods/MSDTrans/map_{data_name}",
    'FIDTM': f"../my_methods/fidtm_mamba_src/map_{data_name}",
}

save_path = './3_visul_results_dronergbt'
if os.path.exists(save_path):
    shutil.rmtree(save_path)
os.makedirs(save_path)



def save_drow_text2map(text, density, save_file):
    # 如果density是tensor，转换为numpy
    if hasattr(density, 'detach'):
        density = density.detach().cpu().numpy()
    
    # 创建图形，调整参数去除白边
    fig, ax = plt.subplots(figsize=(10, 8))
    
    # 显示密度图
    ax.imshow(density, cmap="jet")
    ax.axis('off')  # 关闭坐标轴
    
    # 设置边界为0，去除白边
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
    
    # 文字设置
    font_size = 100
    font_color = 'white'

    # 添加文字（带背景框）
    ax.text(0.5, 0.03, text, 
            fontsize=font_size, 
            color=font_color,
            horizontalalignment='center',
            verticalalignment='bottom',
            transform=ax.transAxes,
    )
    
    # 保存时去除白边
    plt.savefig(save_file, dpi=300, bbox_inches='tight', pad_inches=0)
    plt.close()


results = ['1433', '591', '1590', '701']

for name in results:
    # 1.获取原始数据
    gt_name = f'{name}R.xml'
    et_name = f'{name}_RGB.pt'

    gt_file = os.path.join(gt_path, gt_name)
    rgb_file = os.path.join(rgb_path, gt_name.replace('R.xml', '.jpg'))
    rgb_img = plt.imread(rgb_file)
    image_shape = rgb_img.shape[:2] # (h, w)

    # 2.gt_num, gt_map
    points = get_xml_points(gt_file)
    gt_num = len(points)
    gt_map = get_map_from_points(points, image_shape, sigma=8)
    print(gt_map.shape)
    
    save_name = f"{name}_GT_{int(gt_num)}.png"
    save_file = os.path.join(save_path, save_name)
    text = f"GT: {int(gt_num)}"
    save_drow_text2map(text, gt_map, save_file)

    # 3.获取方法结果
    for i, method in enumerate(method_results.keys()):
        et_file = os.path.join(method_results[method], et_name)

        # map, num
        if method == 'FIDTM':
            et_map, et_num = get_fidtm_map_num(et_file)
        else:
            et_map, et_num = get_map_num(et_file, image_shape)

        print(et_map.shape)
        
        save_name = f"{name}_{method}_{int(et_num)}.png"
        save_file = os.path.join(save_path, save_name)
        text = f"ET: {int(et_num)}"
        save_drow_text2map(text, et_map, save_file)


